import numpy as np
from math import *
def calcEnt(dataset):
    """
    :param dataset:np.array
    :return:
    """
    ydata=dataset[:,-1]
    ydatasize=len(ydata)
    ydataclass={}
    sument=0
    for i in ydata:
        if i not in ydataclass.keys():
            ydataclass[i]=1
        else:
            ydataclass[i]+=1
    for dataclass in ydataclass.keys():
        temp=ydataclass[dataclass]/ydatasize
        sument+=-(temp*log2(temp))
    return sument
def split(dataset,i,value):
    """
    删除第i列的值为value的元素，并把其余的拼接起来
    :param dataset:
    :param i:
    :param value:
    :return:
    """
    spliteddata=[]
    for data in dataset:
        if data[i]==value:
            temp=np.hstack((data[:i],data[i+1:]))
            temp=list(temp)
            spliteddata.append(temp)
    spliteddata=np.array(spliteddata)
    return spliteddata
def chbestspt(dataset):
    datax=dataset[:,:-1]
    baseEnt=calcEnt(dataSet)
    maxgain=0
    maxgaini=0
    for i in range(len(datax[0])):
        xdatai=list(datax[:,i])
        enti=0
        for j in set(xdatai):
            splitdataset=split(dataset,i,j)
            entij=calcEnt(splitdataset)
            enti+=len(splitdataset)/len(dataset)*entij
        gaini=baseEnt-enti
        if gaini>maxgain:
            maxgain=gaini
            maxgaini=i
    return maxgaini
def producetree(dataset,labels):
    xdata = dataset[:, :-1]
    ydata = dataset[:, -1]
    if list(ydata).count(ydata[0])==len(ydata):
        return ydata[0]
    besti=chbestspt(dataset)
    bestilabel=labels[besti]
    mytree={bestilabel:{}}
    del labels[besti]
    for value in set(list(xdata[:,besti])):
        sublabels=labels[:]
        ispliteddata=split(dataset,besti,value)
        mytree[bestilabel][value]=producetree(ispliteddata,sublabels)
    print(mytree)
    return mytree
dataSet = np.loadtxt('data1.csv',dtype=None,delimiter=',')
labels = ['A', 'B', 'C', 'D', 'E', 'F']
tree=producetree(dataSet,labels)